import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from enum import Enum
import random

# 飞行模式枚举
class FlightMode(Enum):
    TAKEOFF = 0
    WAYPOINT = 1
    RETURN_HOME = 2
    LANDING = 3

# 四旋翼参数
MASS = 1.0         # kg
GRAVITY = 9.81     # m/s^2
DT = 0.01          # 时间步长
SIM_TIME = 30      # 仿真时长（秒）
INERTIA = np.array([0.1, 0.1, 0.1])  # 转动惯量

# 飞行参数
TAKEOFF_HEIGHT = 20.0  # 起飞高度
LANDING_SPEED = 0.5   # 降落速度
WAYPOINT_RADIUS = 0.1 # 到达航路点的判定半径
HOME_POSITION = np.array([0.0, 0.0, 0.0])  # 返航点位置

# 航路点参数
MAX_WAYPOINTS = 3     # 每架无人机的最大航路点数
MAX_DISTANCE = 100    # 最大航路点距离

# 速度控制参数
MAX_SPEED = 15.0      # 最大飞行速度 (m/s)
APPROACH_DISTANCE = 2.0  # 接近航路点时开始减速的距离

# 多机参数
NUM_DRONES = 1200      # 无人机数量
DRONE_COLORS = plt.cm.get_cmap('hsv', NUM_DRONES)  # 为每架无人机分配不同颜色

# PID 控制器
class PID:
    def __init__(self, kp, ki, kd, setpoint=0):
        self.kp = kp
        self.ki = ki
        self.kd = kd
        self.setpoint = setpoint
        self.integral = 0
        self.last_error = 0

    def update(self, value, dt):
        error = self.setpoint - value
        self.integral += error * dt
        derivative = (error - self.last_error) / dt
        self.last_error = error
        return self.kp * error + self.ki * self.integral + self.kd * derivative

    def set_setpoint(self, setpoint):
        self.setpoint = setpoint

# 无人机类
class Drone:
    def __init__(self, drone_id):
        self.id = drone_id
        self.color = DRONE_COLORS(drone_id)
        
        # 状态变量
        self.pos = np.array([0.0, 0.0, 0.0])   # 位置 x, y, z
        self.vel = np.array([0.0, 0.0, 0.0])   # 速度 vx, vy, vz
        self.att = np.array([0.0, 0.0, 0.0])   # 姿态角 roll, pitch, yaw
        self.ang_vel = np.array([0.0, 0.0, 0.0])  # 角速度
        
        # 随机生成航路点 (不包括起点)
        num_waypoints = random.randint(1, MAX_WAYPOINTS)
        self.waypoints = []
        for _ in range(num_waypoints):
            x = random.uniform(-MAX_DISTANCE, MAX_DISTANCE)
            y = random.uniform(-MAX_DISTANCE, MAX_DISTANCE)
            z = random.uniform(TAKEOFF_HEIGHT*0.5, TAKEOFF_HEIGHT*1.5)
            self.waypoints.append(np.array([x, y, z]))
        
        # 飞行模式
        self.mode = FlightMode.WAYPOINT
        self.current_waypoint_index = 0
        self.active = True
        
        # PID控制器
        self.pid_x = PID(2.0, 0.0, 1.0)
        self.pid_y = PID(2.0, 0.0, 1.0)
        self.pid_z = PID(10.0, 2.0, 5.0)
        self.pid_vx = PID(2.0, 0.0, 0.5)
        self.pid_vy = PID(2.0, 0.0, 0.5)
        self.pid_vz = PID(2.0, 0.0, 0.5)
        self.pid_roll = PID(2.0, 0.0, 0.5)
        self.pid_pitch = PID(2.0, 0.0, 0.5)
        self.pid_yaw = PID(2.0, 0.0, 0.5)
        
        # 初始化PID设定点
        self.pid_z.set_setpoint(TAKEOFF_HEIGHT)
        
        # 历史数据
        self.positions = []
        self.attitudes = []
        self.velocities = []
        
        # 运动计算相关
        self.prev_vel = np.array([0.0, 0.0, 0.0])
        self.acc = np.array([0.0, 0.0, 0.0])
        self.calculated_att = np.array([0.0, 0.0, 0.0])
        self.filtered_pitch = 0.0
        self.filtered_yaw = 0.0
    
    def calculate_velocity_setpoint(self, target_pos):
        """计算速度设定点，考虑最大速度和接近目标时的减速"""
        # 计算到目标的向量
        direction = target_pos - self.pos
        distance = np.linalg.norm(direction)
        
        if distance < 0.01:  # 防止除以零
            return np.zeros(3)
        
        # 单位方向向量
        direction_unit = direction / distance
        
        # 根据距离计算目标速度
        if distance > APPROACH_DISTANCE:
            # 远距离时使用最大速度
            return direction_unit * MAX_SPEED
        else:
            # 接近目标时线性减速
            speed_factor = distance / APPROACH_DISTANCE
            return direction_unit * MAX_SPEED * speed_factor
    
    def calculate_attitude_from_motion(self):
        """根据速度和加速度计算平滑的姿态信息"""
        MAX_PITCH_ANGLE = 20  # 最大俯仰角(度)
        MAX_ROLL_ANGLE = 30   # 最大滚转角(度)
        
        # 滤波器参数
        alpha = 0.2  # 低通滤波系数 (0.1-0.3)
        max_pitch_rate = 10.0 * DT  # 最大俯仰角变化率 (度/秒)
        max_yaw_rate = 30.0 * DT    # 最大偏航角变化率 (度/秒)
        
        # 1. 计算航向角 (yaw) - 基于飞行方向
        if np.linalg.norm(self.vel[:2]) > 0.1:  # 只有水平速度足够大时才计算
            new_yaw = np.arctan2(self.vel[1], self.vel[0]) * 57.3  # 转换为度
            
            # 处理角度环绕问题 (从-180到180的跳变)
            if abs(new_yaw - self.filtered_yaw) > 180:
                if new_yaw > self.filtered_yaw:
                    new_yaw -= 360
                else:
                    new_yaw += 360
                    
            # 应用变化率限制和低通滤波
            yaw_diff = np.clip(new_yaw - self.filtered_yaw, -max_yaw_rate, max_yaw_rate)
            self.filtered_yaw += yaw_diff
            self.filtered_yaw = alpha * new_yaw + (1 - alpha) * self.filtered_yaw
        else:
            # 保持当前航向
            self.filtered_yaw = self.calculated_att[2]
        
        # 2. 计算俯仰角 (pitch) - 基于加速度方向和速度大小
        # 计算前向加速度分量 (在无人机坐标系中)
        yaw_rad = np.radians(self.filtered_yaw)
        forward_accel = self.acc[0] * np.cos(yaw_rad) + self.acc[1] * np.sin(yaw_rad)
        
        # 计算前向速度分量 (在无人机坐标系中)
        forward_speed = self.vel[0] * np.cos(yaw_rad) + self.vel[1] * np.sin(yaw_rad)
        
        # 决定俯仰方向 - 基于加速度的方向
        pitch_direction = np.sign(forward_accel) if abs(forward_accel) > 0.1 else 0
        
        # 计算俯仰程度 - 基于速度的大小
        speed_factor = np.clip(abs(forward_speed) / MAX_SPEED, 0.0, 1.0)
        
        # 目标俯仰角 = 方向 * 程度 * 最大角度
        target_pitch = pitch_direction * speed_factor * MAX_PITCH_ANGLE
        
        # 应用变化率限制和低通滤波
        pitch_diff = np.clip(target_pitch - self.filtered_pitch, -max_pitch_rate, max_pitch_rate)
        self.filtered_pitch += pitch_diff
        self.filtered_pitch = alpha * target_pitch + (1 - alpha) * self.filtered_pitch
        
        # 3. 计算滚转角 (roll) - 基于转弯时的向心力
        if np.linalg.norm(self.vel[:2]) > 0.5:  # 只有速度足够大时才计算滚转
            # 计算转弯半径和向心加速度
            lateral_accel = -self.acc[0] * np.sin(yaw_rad) + self.acc[1] * np.cos(yaw_rad)
            target_roll = np.clip(np.arctan(lateral_accel / 9.81) * 57.3, -MAX_ROLL_ANGLE, MAX_ROLL_ANGLE)
        else:
            target_roll = 0.0
        
        # 滚转角也需要平滑
        roll_diff = np.clip(target_roll - self.calculated_att[0], -max_pitch_rate, max_pitch_rate)
        filtered_roll = self.calculated_att[0] + roll_diff
        
        # 更新计算姿态 (roll, pitch, yaw)
        self.calculated_att = np.array([
            alpha * filtered_roll + (1 - alpha) * self.calculated_att[0],
            self.filtered_pitch,
            self.filtered_yaw
        ])
        
        return self.calculated_att
    
    def update_flight_mode(self):
        if self.mode == FlightMode.TAKEOFF:
            if self.pos[2] >= TAKEOFF_HEIGHT:
                if len(self.waypoints) > 0:
                    self.mode = FlightMode.WAYPOINT
                    self.current_waypoint_index = 0
                    self.pid_x.set_setpoint(self.waypoints[0][0])
                    self.pid_y.set_setpoint(self.waypoints[0][1])
                    self.pid_z.set_setpoint(self.waypoints[0][2])
                else:
                    self.mode = FlightMode.RETURN_HOME
                    self.pid_x.set_setpoint(HOME_POSITION[0])
                    self.pid_y.set_setpoint(HOME_POSITION[1])
                    self.pid_z.set_setpoint(HOME_POSITION[2])
        
        elif self.mode == FlightMode.WAYPOINT:
            current_waypoint = self.waypoints[self.current_waypoint_index]
            distance = np.linalg.norm(self.pos - current_waypoint)
            
            if distance < WAYPOINT_RADIUS:
                self.current_waypoint_index += 1
                if self.current_waypoint_index >= len(self.waypoints):
                    self.mode = FlightMode.RETURN_HOME
                    self.pid_x.set_setpoint(HOME_POSITION[0])
                    self.pid_y.set_setpoint(HOME_POSITION[1])
                    self.pid_z.set_setpoint(HOME_POSITION[2])
                else:
                    next_waypoint = self.waypoints[self.current_waypoint_index]
                    self.pid_x.set_setpoint(next_waypoint[0])
                    self.pid_y.set_setpoint(next_waypoint[1])
                    self.pid_z.set_setpoint(next_waypoint[2])
        
        elif self.mode == FlightMode.RETURN_HOME:
            distance = np.linalg.norm(self.pos - HOME_POSITION)
            if distance < WAYPOINT_RADIUS:
                self.mode = FlightMode.LANDING
                self.pid_z.set_setpoint(0)
        
        elif self.mode == FlightMode.LANDING:
            # 让 setpoint 以最大下降速度递减
            new_setpoint = max(self.pid_z.setpoint - LANDING_SPEED * DT, 0)
            self.pid_z.set_setpoint(new_setpoint)
            if self.pos[2] <= 0.1:
                self.active = False
    
    def update(self):
        if not self.active:
            return
        
        # 随机风的影响
        wind = np.random.normal(0, 0.0001, size=3)  # 轻微的随机风
        self.vel += wind
        
        # 更新飞行模式
        self.update_flight_mode()

        # 计算速度设定点
        if self.mode == FlightMode.TAKEOFF:
            vx_setpoint = 0
            vy_setpoint = 0
            vz_setpoint = min(MAX_SPEED, (TAKEOFF_HEIGHT - self.pos[2]) / 2)
        elif self.mode == FlightMode.WAYPOINT:
            target_pos = self.waypoints[self.current_waypoint_index]
            velocity_setpoint = self.calculate_velocity_setpoint(target_pos)
            vx_setpoint, vy_setpoint, vz_setpoint = velocity_setpoint
        elif self.mode == FlightMode.RETURN_HOME:
            velocity_setpoint = self.calculate_velocity_setpoint(HOME_POSITION)
            vx_setpoint, vy_setpoint, vz_setpoint = velocity_setpoint
        elif self.mode == FlightMode.LANDING:
            vx_setpoint = 0
            vy_setpoint = 0
            vz_setpoint = -min(LANDING_SPEED, (self.pos[2] - self.pid_z.setpoint) / 2)
        
        # 设置速度PID控制器的设定点
        self.pid_vx.set_setpoint(vx_setpoint)
        self.pid_vy.set_setpoint(vy_setpoint)
        self.pid_vz.set_setpoint(vz_setpoint)
        
        # 速度PID控制器输出
        fx = self.pid_vx.update(self.vel[0], DT)
        fy = self.pid_vy.update(self.vel[1], DT)
        fz = self.pid_vz.update(self.vel[2], DT) + MASS * GRAVITY

        # 计算姿态控制力矩
        tau_roll = self.pid_roll.update(self.att[0], DT)
        tau_pitch = self.pid_pitch.update(self.att[1], DT)
        tau_yaw = self.pid_yaw.update(self.att[2], DT)

        # 合力
        force = np.array([fx, fy, fz])
        torque = np.array([tau_roll, tau_pitch, tau_yaw])

        # 加速度
        self.acc = force / MASS - np.array([0, 0, GRAVITY])
        ang_acc = torque / INERTIA

        # 更新速度和位置
        self.vel += self.acc * DT
        self.pos += self.vel * DT
        
        # 防止降到地面以下
        if self.pos[2] < 0:
            self.pos[2] = 0
            self.vel[2] = 0

        # 更新角速度和姿态
        self.ang_vel += ang_acc * DT
        self.att += self.ang_vel * DT

        # 计算基于运动的姿态信息
        self.calculated_att = self.calculate_attitude_from_motion()
        
        # 记录历史数据
        self.positions.append(self.pos.copy())
        self.attitudes.append(self.calculated_att.copy())
        self.velocities.append(self.vel.copy())

# 创建无人机群
drones = [Drone(i) for i in range(NUM_DRONES)]

# 创建图形
fig = plt.figure(figsize=(12, 10))
ax = fig.add_subplot(111, projection='3d')

# 设置坐标轴范围
ax.set_xlim([-MAX_DISTANCE*1.1, MAX_DISTANCE*1.1])
ax.set_ylim([-MAX_DISTANCE*1.1, MAX_DISTANCE*1.1])
ax.set_zlim([0, MAX_DISTANCE*1.1])
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')

# 绘制返航点
ax.scatter(HOME_POSITION[0], HOME_POSITION[1], HOME_POSITION[2], color='red', label='Home', s=100)


# 为每架无人机创建图形对象
drone_objects = []
for drone in drones:
    # 轨迹线
    traj, = ax.plot([], [], [], '-', color=drone.color, alpha=0.3, linewidth=1)
    
    # 当前位置点
    point, = ax.plot([], [], [], 'o', color=drone.color, markersize=3)
    
    # 速度向量（箭头简化为线段）
    vec, = ax.plot([], [], [], '-', color=drone.color, linewidth=1)
    
    drone_objects.append({
        'trajectory': traj,
        'point': point,
        'velocity_vec': vec
    })

def update(frame):
    active_drones = 0
    for i, drone in enumerate(drones):
        if drone.active:
            drone.update()
            active_drones += 1

        objs = drone_objects[i]
        
        # 更新轨迹线
        pos_arr = np.array(drone.positions)
        if len(pos_arr) > 1:
            objs['trajectory'].set_data(pos_arr[:,0], pos_arr[:,1])
            objs['trajectory'].set_3d_properties(pos_arr[:,2])
        else:
            objs['trajectory'].set_data([], [])
            objs['trajectory'].set_3d_properties([])
        
        # 更新当前位置点
        objs['point'].set_data([drone.pos[0]], [drone.pos[1]])
        objs['point'].set_3d_properties([drone.pos[2]])
        
        # 更新速度向量
        tail = drone.pos
        head = drone.pos + drone.vel * 0.2  # 缩放速度向量
        objs['velocity_vec'].set_data([tail[0], head[0]], [tail[1], head[1]])
        objs['velocity_vec'].set_3d_properties([tail[2], head[2]])
    
    # 可以在绘图区显示剩余活跃无人机数
    ax.set_title(f"Frame {frame}, Active drones: {active_drones}/{NUM_DRONES}")
    
    return [obj for d in drone_objects for obj in d.values()]

# 创建动画
ani = FuncAnimation(fig, update, frames=int(SIM_TIME/DT), interval=DT*1000, blit=True)

plt.legend()
plt.tight_layout()
plt.show()
